import sys
import os
import zipfile
from PyQt5.QtWidgets import QApplication, QMainWindow, QFileDialog, QMessageBox, QDialog, QVBoxLayout, QLabel, QLineEdit, QComboBox, QDialogButtonBox
from main import Ui_MainWindow
import docx
from openpyxl import load_workbook
import shutil
import traceback

class ImageSettingsDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle('图片设置')

        layout = QVBoxLayout(self)

        self.lbl_width = QLabel('宽度:', self)
        self.edit_width = QLineEdit(self)
        self.edit_width.setText("100")
        layout.addWidget(self.lbl_width)
        layout.addWidget(self.edit_width)

        self.lbl_height = QLabel('高度:', self)
        self.edit_height = QLineEdit(self)
        self.edit_height.setText("100")
        layout.addWidget(self.lbl_height)
        layout.addWidget(self.edit_height)

        self.lbl_alignment = QLabel('对齐方式:', self)
        self.combo_alignment = QComboBox(self)
        self.combo_alignment.addItems(['左对齐', '居中', '右对齐'])
        layout.addWidget(self.lbl_alignment)
        layout.addWidget(self.combo_alignment)

        button_box = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel, self)
        button_box.accepted.connect(self.accept)
        button_box.rejected.connect(self.reject)
        layout.addWidget(button_box)

class MainWindow(QMainWindow, Ui_MainWindow):
    def __init__(self, parent=None):
        super(MainWindow, self).__init__(parent)
        self.setupUi(self)
        #uic.loadUi('main.ui', self)  # 加载 UI 文件

        self.word_template = None
        self.excel_data = None
        self.keys = None

        # 将按钮连接到对应的功能
        self.btn_select_word_template.clicked.connect(self.select_word_template)
        self.btn_select_excel_data.clicked.connect(self.select_excel_data)
        self.btn_generate_word_files.clicked.connect(self.generate_word_files)
        self.btn_batch_zip.clicked.connect(self.batch_zip)

    # 选择 Word 模板
    def select_word_template(self):
        options = QFileDialog.Options()
        options |= QFileDialog.ReadOnly
        self.word_template, _ = QFileDialog.getOpenFileName(self, '选择 Word 模板', '', 'Word 文档 (*.docx)', options=options)
        if self.word_template:
            self.lbl_word_template.setText(os.path.basename(self.word_template))

    # 选择 Excel 数据表
    def select_excel_data(self):
        options = QFileDialog.Options()
        options |= QFileDialog.ReadOnly
        self.excel_data, _ = QFileDialog.getOpenFileName(self, '选择 Excel 数据表', '', 'Excel 文档 (*.xlsx)', options=options)
        if self.excel_data:
            self.lbl_excel_data.setText(os.path.basename(self.excel_data))

    # 生成 Word 文件
    def generate_word_files(self):
        if not self.word_template or not self.excel_data:
            QMessageBox.warning(self, '警告', '请先选择 Word 模板和 Excel 数据表')
            return

        try:
            # 从 Excel 中提取图片
            # images_path = 'imagesout'
            # if os.path.exists(images_path):
            #     shutil.rmtree(images_path)
            # os.makedirs(images_path, exist_ok=True)

            # with zipfile.ZipFile(self.excel_data, 'r') as zf:
            #     for file in zf.namelist():
            #         if file.startswith('xl/media/'):
            #             zf.extract(file, images_path)

            # 读取 Excel 数据
            wb = load_workbook(self.excel_data, data_only=True)
            ws = wb.active
            self.keys = [cell.value for cell in ws[1]]

            # 设置图片大小
            # image_settings_dialog = ImageSettingsDialog(self)
            # result = image_settings_dialog.exec_()
            # if result == QDialog.Accepted:
            #     width = int(image_settings_dialog.edit_width.text())
            #     height = int(image_settings_dialog.edit_height.text())
            #     alignment = image_settings_dialog.combo_alignment.currentIndex()

            # 生成 Word 文件
            for idx, row in enumerate(ws.iter_rows(min_row=2, values_only=True)):
                data = dict(zip(self.keys, row))
                # if '图片' in self.keys:  # 如果存在图片列，将图片路径添加到数据中
                #     img_filename = "image" + str(idx + 1)  # 根据对应行数生成图片文件名
                #     img_path = os.path.join(images_path, "xl", "media", img_filename).replace('\\', '/')
                #     if os.path.exists(img_path + ".png"):
                #         img_path += ".png"
                #     elif os.path.exists(img_path + ".jpg"):
                #         img_path += ".jpg"
                #     elif os.path.exists(img_path + ".jpeg"):
                #         img_path += ".jpeg"
                #     elif os.path.exists(img_path + ".gif"):
                #         img_path += ".gif"
                #     else:
                #         img_path = None
                #
                #     if img_path:
                #         data['图片'] = img_path
                #     else:
                #         data.pop('图片')  # 如果该行没有图片，将图片键从数据中删除
                # self.generate_word_file(data, width, height, alignment)
                self.generate_word_file(data)

            QMessageBox.information(self, '完成', 'Word 文件已生成')
        except Exception as e:
            exc_type, exc_value, exc_traceback = sys.exc_info()
            traceback_str = ''.join(traceback.format_exception(exc_type, exc_value, exc_traceback))
            with open('log.txt', 'a') as log_file:
                log_file.write('generate_word_files error: ' + traceback_str + '\n')
            QMessageBox.critical(self, '错误', '生成文件出错，请查看 log.txt 获取详细信息')

    # 生成单个 Word 文件
    # def generate_word_file(self, data, width, height, alignment):
    def generate_word_file(self, data):
        doc = docx.Document(self.word_template)

        # 处理插入图片和替换文本的逻辑
        for paragraph in doc.paragraphs:
            for run in paragraph.runs:
                for key, value in data.items():
                    if '{{' + key + '}}' in run.text or '｛｛' + key + '｝｝' in run.text:
                        # if key == self.keys[-1]:  # 如果是图片的占位符
                        #     try:
                        #         run.clear()
                        #         run.add_picture(value, width=docx.shared.Pt(width), height=docx.shared.Pt(height))
                        #         if alignment == 0:
                        #             paragraph.alignment = docx.enum.text.WD_ALIGN_PARAGRAPH.LEFT
                        #         elif alignment == 1:
                        #             paragraph.alignment = docx.enum.text.WD_ALIGN_PARAGRAPH.CENTER
                        #         else:
                        #             paragraph.alignment = docx.enum.text.WD_ALIGN_PARAGRAPH.RIGHT
                        #     except Exception as e:
                        #         exc_type, exc_value, exc_traceback = sys.exc_info()
                        #         traceback_str = ''.join(traceback.format_exception(exc_type, exc_value, exc_traceback))
                        #         with open('log.txt', 'a') as log_file:
                        #             log_file.write('generate_word_file error: ' + traceback_str + '\n')
                        #         QMessageBox.critical(self, '错误', '插入图片失败，请查看 log.txt 获取详细信息')
                        # else:
                        run.text = run.text.replace('{{' + key + '}}', str(value)).replace('｛｛' + key + '｝｝', str(value))
        desktop = os.path.expanduser("./outdocx")
        filename = '{}.docx'.format(data[list(data.keys())[0]])
        filename = os.path.join(desktop, '{}.docx'.format(data[list(data.keys())[0]]))
        os.makedirs(desktop, exist_ok=True)
        doc.save(filename)

    # 批量压缩 Word 文件
    def batch_zip(self):
        if not os.path.exists('./outdocx'):
            QMessageBox.warning(self, '警告', '请先生成 Word 文件')
            return

        with zipfile.ZipFile('output.zip', 'w') as zipf:
            for root, _, files in os.walk('./outdocx'):
                for file in files:
                    zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), './outdocx'))

        QMessageBox.information(self, '完成', '已将生成的 Word 文件压缩为 output.zip')

if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())
